from langchain_deepseek import ChatDeepSeek

from src.nl2sql.utils import extract_configuration_tables, get_tables_schema
from src.nl2sql.state import Nl2SqlState, QueryOutput, TableOutput
from src.nl2sql.prompts import system, choice_table_prompt, table_schema, schema, second_schema
from langchain_core.prompts import ChatPromptTemplate
from src.chat_utils import create_chat, query_mysql_pk, create_db


# 开始节点，加载数据库链接
def start_entry(state: Nl2SqlState):
    db = create_db()
    return {"db": db}


# 根据问题查找主业务表
def retrieve_main_table(state: Nl2SqlState):
    db = state["db"]
    llm = create_chat()
    res = query_mysql_pk(state["nl2sql_messages"][-1].content)[0]
    table_names = [r["entity"]["model"] for r in res][0]
    if table_names == 'mysql,itsdb,se_life_cycle':
        table_names = 'mysql,itsdb,ding_talk_attendance'
    # db.get_usable_table_names()
    prompt = choice_table_prompt.format(table_names=table_names,
                                  dialect=db.dialect)
    query_prompt_template = ChatPromptTemplate.from_messages([
        ("system", prompt),
        ("user", "Question:{input}")
    ])
    prompt = query_prompt_template.invoke(
        {
            "input": state["nl2sql_messages"],
        }
    )
    structured_llm = llm.with_structured_output(TableOutput)
    main_table = structured_llm.invoke(prompt)
    # 后续更新表注解之后，这段代码需要修改
    main_table = main_table["table_name"].split(",")[-1]
    table_schema = get_tables_schema(db, main_table, llm)
    # print(table_schema)
    return {"main_table": main_table, "table_schema": table_schema}



# 根据主表标注信息查找配置表
def choice_configuration_tables(state: Nl2SqlState):

    schema = state["table_schema"]
    configuration_tables = extract_configuration_tables(schema)
    configuration_tables_schema = list()
    description = list()
    if configuration_tables:
        db = state["db"]
        llm = create_chat()
        ts = get_tables_schema(db, configuration_tables, llm)
        configuration_tables_schema.extend(ts)
        description.extend([state["main_table"] + "表跟" + table + "表关联" for table in configuration_tables])
    return {"configuration_tables": configuration_tables,
            "configuration_tables_schema": configuration_tables_schema,
            "description": description}



def generate_sql(state: Nl2SqlState):
    # llm = create_chat(ollama_model="qwen2.5-coder:14b")
    llm = ChatDeepSeek(
        model="deepseek-chat",
        temperature=0,
        max_tokens=None,
        timeout=None,
        max_retries=2,
        api_key="sk-d58010675dde402a92107369a296bfd8",
        # other params...
    )
    db = state["db"]
    table_schema = [state["table_schema"]] if  isinstance(state["table_schema"], str)\
        else state["table_schema"]
    configuration_tables_schema = state["configuration_tables_schema"]
    table_names = table_schema + configuration_tables_schema
    system_prompt = system.format(table_names=table_names,
                    dialect=db.dialect,
                    description="".join(state["description"]))
    query_prompt_template = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("user", "Question:{input}")
    ])
    prompt = query_prompt_template.invoke(
        {
            "input": state["nl2sql_messages"],
        }
    )
    structured_llm = llm.with_structured_output(QueryOutput)
    result = structured_llm.invoke(prompt)
    return {"sql": result["sql"]}

#判断是否含有第二层链接表
def should_find_tables(state: Nl2SqlState):

    flag = "generate_sql"
    configuration_tables_schema = state["configuration_tables_schema"]
    for schema in configuration_tables_schema:
        tables = extract_configuration_tables(schema)
        if tables:
            flag = "continue_find"
            break
    return flag

def continue_find(state: Nl2SqlState):
    schemas = state["configuration_tables_schema"]
    configuration_tables_schema = list()
    configuration_tables_list = list()
    description = list()
    for index, schema in enumerate(schemas):
        configuration_tables = extract_configuration_tables(schema)
        configuration_tables_list.extend(configuration_tables)
        if configuration_tables:
            db = state["db"]
            llm = create_chat()
            configuration_tables_schema.extend(get_tables_schema(db, configuration_tables, llm))
            description.extend([state["configuration_tables"][index] + "表跟" + table + "表关联" for table in configuration_tables])
    state["configuration_tables"].extend(configuration_tables_list)
    state["configuration_tables_schema"].extend(configuration_tables_schema)
    state["description"].extend(description)
    # return {"configuration_tables": state["configuration_tables"].extend(configuration_tables_list),
    #         "configuration_tables_schema": state["configuration_tables_schema"].extend(configuration_tables_schema)}
    return None




if __name__ == '__main__':
    # query = "查询2025年3月15日和2025年3月17日三江化工和嘉兴华雯化工上游来水的pH和cod数据"
    # query = "查询2025年3月15日到2025年3月17日三江化工和嘉兴华雯化工上游来水以及下游排海的ph,cod数据"
    # query = "查询2025年3月10日到2025年3月16日上游来水的数据"
    # query = "查询2025年3月10日到2025年3月16日下游排海的数据"
    # query = "查询2025年3月9日到2025年3月12日厂内在线数据"
    # query = "查询2025年3月9日到11日实验室检测二沉池(A)、二沉池(B)和磁混凝的异常数据"
    # query = "查询厂内运行巡检打卡记录数据"
    # query = "查询运行缺卡数据"
    # query = "查询2025年3月10日到2025年3月16日的2#机主和深度处理A的用电情况"
    # query = "查询2025年3月10日到2025年3月16日的生活用水量"
    # query = "查询2025年1月1日到2025年3月16日的用蒸汽情况"
    # query = "查询2025年3月10日到3月19日营养液、液氧的药剂出入库情况"
    # query = "查询2025年3月10日到3月19日营养液、液氧的药剂总量情况"
    # query = "查询药剂液氮入库情况"
    # query = "入库时间为2025年3月污泥外运平湖石化情况"
    # query = "查询2025年3月20日运行打卡数据"
    # query = "查询2025年3月16日到23日运行缺卡数据"
    # query = "查询考勤数据"
    # query = "查询2025年3月考勤打卡数据"
    # query = "获取三江化工和传化合成的今天上游来水和下游排海"
    querys = ["查询2025年3月15日和2025年3月17日三江化工和嘉兴华雯化工上游来水的pH和cod数据",

                "查询2025年3月15日到2025年3月17日三江化工和嘉兴华雯化工上游来水以及下游排海的ph,cod数据",
                "查询2025年3月10日到2025年3月16日上游来水的数据",
                "查询2025年3月10日到2025年3月16日下游排海的数据",
                "查询2025年3月9日到2025年3月12日厂内在线数据",
                "查询2025年3月9日到11日实验室检测二沉池(A)、二沉池(B)和磁混凝的异常数据",
                "查询厂内运行巡检打卡记录数据",
                "查询运行缺卡数据",
                "查询2025年3月10日到2025年3月16日的2#机主和深度处理A的用电情况",
                "查询2025年3月10日到2025年3月16日的生活用水量",
                "查询2025年1月1日到2025年3月16日的用蒸汽情况",
                "查询2025年3月10日到3月19日营养液、液氧的药剂出入库情况",
                "查询2025年3月10日到3月19日营养液、液氧的药剂总量情况",
                "查询药剂液氮入库情况",
                "入库时间为2025年3月污泥外运平湖石化情况",
                "查询2025年3月20日运行打卡数据",
                "查询2025年3月16日到23日运行缺卡数据",
                "查询考勤数据",
                "查询2025年3月考勤打卡数据",
                "获取三江化工和传化合成的今天上游来水和下游排海"]
    for query in querys[0:1]:
        print(query+ "sql: ",retrieve_main_table({"nl2sql_messages": [{"role": "system", "content": "system_prompt"},
                                                                      {"role": "user", "content": query}]}).get("table_name"))
    # generate_sql({"nl2sql_messages": [{"role": "user", "content": "查询2025年3月9日到2025年3月12日厂内在线数据"}]})